import os
import torch
import signal, subprocess, submitit
import torch.backends.cudnn as cudnn
import torch.distributed as dist

from pathlib import Path
from ecgcmr.utils.misc import fix_seed


def init_dist_node(args):
    args.ngpus_per_node = torch.cuda.device_count()

    # requeue job on SLURM preemption
    signal.signal(signal.SIGUSR1, handle_sigusr1)
    signal.signal(signal.SIGTERM, handle_sigterm)

    # find a common host name on all nodes
    cmd = 'scontrol show hostnames ' + os.getenv('SLURM_JOB_NODELIST')
    stdout = subprocess.check_output(cmd.split())
    host_name = stdout.decode().splitlines()[0]
    print(host_name, args.port)
    args.dist_url = f'tcp://{host_name}:{args.port}'

    # distributed parameters
    args.rank = int(os.getenv('SLURM_NODEID')) * args.ngpus_per_node
    args.world_size = int(os.getenv('SLURM_NNODES')) * args.ngpus_per_node


# Signal Handlers
def handle_sigusr1(signum, frame):
    os.system(f'scontrol requeue {os.getenv("SLURM_JOB_ID")}')
    exit()

def handle_sigterm(signum, frame):
    pass

def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True

def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print

def init_dist_gpu(gpu, args):

    if args.slurm:
        job_env = submitit.JobEnvironment()

        args.output_dir = Path(str(args.output_dir).replace("%j", str(job_env.job_id)))
        
        args.gpu = job_env.local_rank
        args.rank = job_env.global_rank
        args.world_size = job_env.num_tasks

        # Set necessary environment variables for PyTorch distributed operations
        os.environ['LOCAL_RANK'] = str(args.gpu)
        os.environ['RANK'] = str(args.rank)
        os.environ['WORLD_SIZE'] = str(args.world_size)

    else:
        args.gpu = gpu
        args.rank += gpu

    dist.init_process_group(backend='gloo', init_method=args.dist_url,
                            world_size=args.world_size, rank=args.rank)
    
    print("Backend is loaded", is_dist_avail_and_initialized())

    fix_seed()

    torch.cuda.set_device(args.gpu)
    cudnn.benchmark = True
    dist.barrier()

    args.main = (args.rank == 0)
    setup_for_distributed(args.main)